#!/usr/bin/env bash

set -Eeuxo pipefail

usage() {
    cat <<EOF
Usage: $0 --network <name> --fork-config <fork-config.json> --working-dir <dir> [options]

Required:
  --network                   Name of the network
  --fork-config               Path to fork_config.json
  --working-dir               Working directory

Optional:
  --precomputed-block-prefix  Precomputed block prefix

This script is used to validate that an installed package is correct
according to an exported fork_config.json file.

Inputs:
- The exported mainnet full config.json fork config, with all accounts
- A working directory where ledgers/configs will be created
- Installed MINA_EXE (default: mina) and MINA_GENESIS_EXE (default: mina-create-genesis) programs
- Installed MINA_LEGACY_GENESIS_EXE program (default: mina-legacy-create-genesis)
- PACKAGED_DAEMON_CONFIG (default: /var/lib/coda/config_*.json)
  the runtime config generated by the HF packaging
- CREATE_RUNTIME_CONFIG (default: mina-hf-create-runtime-config)
  a copy of ./scripts/hardfork/create_runtime_config.sh
- GENESIS_LEDGER_DIR (default: /var/lib/coda)
  where the tarfiles to verify are stored
- FORKING_FROM_CONFIG_JSON (default: /var/lib/coda/mainnet.json)
  the pre-fork genesis ledger
- SECONDS_PER_SLOT (default: 180)
- PRECOMPUTED_FORK_BLOCK (default: fetches with gsutil)
- GSUTIL (default: gsutil)
  the Google Cloud Storage utility if the PRECOMPUTED_FORK_BLOCK isn't a file

Ensures:
- The accounts listed in config.json are the ones in the PACKAGED_DAEMON_CONFIG
- The genesis ledger directory tarfile contents match reference copies created
  fresh from config.json. This takes a long time (>20min) due to rehashing.

Outputs:
- Exit code 0 if validated, 1 otherwise.
EOF

    exit 1
}

PRECOMPUTED_BLOCK_PREFIX=""

# Parse named flags
while [[ $# -gt 0 ]]; do
    case $1 in
        --network)
            NETWORK_NAME="$2"
            shift 2
            ;;
        --fork-config)
            FORK_CONFIG="$2"
            shift 2
            ;;
        --working-dir)
            WORKDIR="$2"
            shift 2
            ;;
        --precomputed-block-prefix)
            PRECOMPUTED_BLOCK_PREFIX="$2"
            shift 2
            ;;
        -*)
            echo "Unknown option $1"
            usage
            ;;
        *)
            echo "Unexpected argument $1"
            usage
            ;;
    esac
done

# Validate required arguments
if [[ -z "$NETWORK_NAME" ]]; then
    echo "Error: --network is required"
    usage
fi

if [[ -z "$FORK_CONFIG" ]]; then
    echo "Error: --fork-config is required"
    usage
fi

if [[ -z "$WORKDIR" ]]; then
    echo "Error: --working-dir is required"
    usage
fi

source_build_fallback() {
    v=$(command -v "$1")
    if [ ! -x "$v" ]; then
        if [ ! -x "$2" ]; then
            echo "Error: program not found in PATH (as $1) or relative to cwd (as $2)" >&2
            exit 1
        else
            echo "$2"
        fi
    else
      echo "$1"
    fi
}

MINA_EXE=${MINA_EXE:-$(source_build_fallback mina ./_build/default/src/app/cli/src/mina.exe)}
MINA_GENESIS_EXE=${MINA_GENESIS_EXE:-$(source_build_fallback mina-create-genesis ./_build/default/src/app/runtime_genesis_ledger/runtime_genesis_ledger.exe)}
MINA_LEGACY_GENESIS_EXE=${MINA_LEGACY_GENESIS_EXE:-$(source_build_fallback mina-create-legacy-genesis ./runtime_genesis_ledger_of_mainnet.exe)}
CREATE_RUNTIME_CONFIG=${CREATE_RUNTIME_CONFIG:-$(source_build_fallback mina-hf-create-runtime-config ./scripts/hardfork/create_runtime_config.sh)}
MINA_ROCKSDB_SCANNER=${MINA_ROCKSDB_SCANNER:-$(source_build_fallback mina-rocksdb-scanner ./_build/default/src/app/rocksdb-scanner/rocksdb_scanner.exe)}
GSUTIL=${GSUTIL:-$(command -v gsutil || echo "")}

if [[ ! -e "${PRECOMPUTED_FORK_BLOCK:-}" && ( -z "$GSUTIL" || ! -x "$GSUTIL" ) ]]; then
    echo "Error: gsutil is required when PRECOMPUTED_FORK_BLOCK is nonexistent path"
    exit 1
fi

installed_config=$(echo /var/lib/coda/config_*.json)
PACKAGED_DAEMON_CONFIG=${PACKAGED_DAEMON_CONFIG:-$installed_config}
if [ ! -e "$PACKAGED_DAEMON_CONFIG" ]; then
    echo "Error: set PACKAGED_DAEMON_CONFIG to the path to the JSON file to verify"
    exit 1
fi


GENESIS_LEDGER_DIR=${GENESIS_LEDGER_DIR:-"/var/lib/coda"}

SECONDS_PER_SLOT=${SECONDS_PER_SLOT:=180}

export FORKING_FROM_CONFIG_JSON=${FORKING_FROM_CONFIG_JSON:-$(source_build_fallback /var/lib/coda/mainnet.json genesis_ledgers/mainnet.json)}

export MINA_LIBP2P_PASS=''

mkdir -p "$WORKDIR"/{ledgers{,-backup},keys}
chmod 700 "$WORKDIR/keys"

fork_block_state_hash=$(jq -r '.proof.fork.state_hash' "$FORK_CONFIG")
fork_block_length=$(jq -r '.proof.fork.blockchain_length' "$FORK_CONFIG")

# Fetch the precomputed fork block from local file or gcloud
fetch_precomputed_fork_block() {
    local network_name=$1
    local fork_block_length=$2
    local fork_block_state_hash=$3
    local precomputed_block_prefix=$4
    local dst=$5

    # Put the fork block where we want it, fetch it from gcloud if necessary
    if [[ -z "${PRECOMPUTED_FORK_BLOCK:-}" ]]; then
        PRECOMPUTED_FORK_BLOCK="${precomputed_block_prefix:-gs://mina_network_block_data/$network_name}-$fork_block_length-$fork_block_state_hash.json"
    fi

    local src="$PRECOMPUTED_FORK_BLOCK"
    if [[ -e "$src" ]]; then
        cp "$src" "$dst"
    else
        "$GSUTIL" cp "$src" "$dst"
    fi
}

fetch_precomputed_fork_block "$NETWORK_NAME" "$fork_block_length" "$fork_block_state_hash" "$PRECOMPUTED_BLOCK_PREFIX" "$WORKDIR/precomputed_fork_block.json"

if [ ! -e "$WORKDIR/keys/p2p" ]; then
    "$MINA_EXE" libp2p generate-keypair --privkey-path "$WORKDIR/keys/p2p"
fi

echo "generating genesis ledgers ... (this may take a while)" >&2

"$MINA_LEGACY_GENESIS_EXE" --config-file "$FORK_CONFIG" --genesis-dir "$WORKDIR/legacy_ledgers" --hash-output-file "$WORKDIR/legacy_hashes.json"

result=$(jq --slurpfile block "$WORKDIR/precomputed_fork_block.json" \
  --slurpfile legacy_hashes "$WORKDIR/legacy_hashes.json" -n '
  ($legacy_hashes[0].epoch_data.staking.hash == $block[0].data.protocol_state.body.consensus_state.staking_epoch_data.ledger.hash and
   $legacy_hashes[0].epoch_data.next.hash == $block[0].data.protocol_state.body.consensus_state.next_epoch_data.ledger.hash and
   $legacy_hashes[0].ledger.hash == $block[0].data.protocol_state.body.blockchain_state.staged_ledger_hash.non_snark.ledger_hash)')

if [ "$result" != "true" ]; then
    echo "Hashes in config $FORK_CONFIG don't match hashes from the precomputed block $PRECOMPUTED_FORK_BLOCK" >&2
    exit 1
fi

"$MINA_GENESIS_EXE" --config-file "$FORK_CONFIG" --genesis-dir "$WORKDIR/ledgers" --hash-output-file "$WORKDIR/hashes.json"

FORK_CONFIG_JSON="$FORK_CONFIG" \
LEDGER_HASHES_JSON="$WORKDIR/hashes.json" \
GENESIS_TIMESTAMP=$(jq -r '.genesis.genesis_state_timestamp' "$PACKAGED_DAEMON_CONFIG") \
"$CREATE_RUNTIME_CONFIG" > "$WORKDIR/config-substituted.json"

echo "exporting ledgers from running node ... (this may take a while)" >&2

override_file="$WORKDIR/override-genesis-timestamp.json"
echo "{\"genesis\":{\"genesis_state_timestamp\":\"$(date -u +"%Y-%m-%dT%H:%M:%SZ")\"}}" > "$override_file"

# export the packaged ledgers in a way where we know which one is which
function extract_ledgers() {
    config_file=$1
    ledger_dir=$2
    json_prefix=$3
    
    "$MINA_EXE" daemon --libp2p-keypair "$WORKDIR/keys/p2p" \
      --config-file "$config_file" --config-file "$override_file" \
      --seed --genesis-ledger-dir "$ledger_dir" \
      --log-level "${MINA_LOG_LEVEL:-info}" &

    while ! "$MINA_EXE" ledger export staged-ledger  | jq . 2>/dev/null >"$json_prefix-staged.json"; do
        sleep 1m
        if ! grep -qFx "$(cat ~/.mina-config/.mina-lock)" <(jobs -rp); then
            echo "daemon died before exporting ledgers" >&2
            exit 1
        fi
    done

    "$MINA_EXE" ledger export staking-epoch-ledger | jq . > "$json_prefix-staking.json"
    "$MINA_EXE" ledger export next-epoch-ledger | jq . > "$json_prefix-next.json"

    "$MINA_EXE" client stop
}

mkdir -p "$WORKDIR/ledgers-downloaded"

mv -t "$WORKDIR/ledgers-backup" "$GENESIS_LEDGER_DIR"/*.tar.gz
if [ -z ${NO_TEST_LEDGER_DOWNLOAD+z} ]; then
    # We removed all *.tar.gz from GENESIS_LEDGER_DIR (in case it's one of default paths)
    # and expect Mina node to be able to download these from S3
    extract_ledgers "$PACKAGED_DAEMON_CONFIG" "$WORKDIR/ledgers-downloaded" "$WORKDIR/downloaded"
    rm -Rf /tmp/coda_cache_dir/*.tar.gz /tmp/s3_cache_dir/*.tar.gz
fi
extract_ledgers "$WORKDIR/config-substituted.json" "$WORKDIR/ledgers" "$WORKDIR/reference"
mv -t "$GENESIS_LEDGER_DIR" "$WORKDIR/ledgers-backup"/*
extract_ledgers "$PACKAGED_DAEMON_CONFIG" "$GENESIS_LEDGER_DIR" "$WORKDIR/packaged"

echo "Performing final comparisons..." >&2
error=0

result=$(jq --slurpfile a "$WORKDIR/config-substituted.json" --slurpfile b "$PACKAGED_DAEMON_CONFIG" -n '
  ($a[0].epoch_data.staking.hash == $b[0].epoch_data.staking.hash and
   $a[0].epoch_data.next.hash == $b[0].epoch_data.next.hash and
   $a[0].ledger.hash == $b[0].ledger.hash)')

if [ "$result" != "true" ]; then
    echo "Packaged config hashes in $PACKAGED_DAEMON_CONFIG not expected compared to $WORKDIR/config-substituted.json" >&2
    error=1
fi

for file in "$WORKDIR"/packaged-*.json; do
    name=$(basename "$file")
    name=${name%.json}
    name=${name#packaged-}

    if ! cmp "$file" "$WORKDIR/reference-$name.json"; then
        echo "Error: $file does not match reference" >&2
        error=1
    fi
    if [ -z ${NO_TEST_LEDGER_DOWNLOAD+z} ]; then
        if ! cmp "$file" "$WORKDIR/downloaded-$name.json"; then
            echo "Error: $file does not match downloaded" >&2
            error=1
        fi
    fi
done

for file in "$WORKDIR"/ledgers/*.tar.gz; do
    tarname=$(basename "$file")
    tarname=${tarname%.tar.gz}
    tardir="$WORKDIR/ledgers/$tarname"
    mkdir -p "$tardir"/{packaged,generated,web}
    tar -xzf "$file" -C "$tardir/generated"
    tar -xzf "$GENESIS_LEDGER_DIR/$tarname.tar.gz" -C "$tardir/packaged"

    base_s3_url="${MINA_LEDGER_S3_BUCKET:-https://s3-us-west-2.amazonaws.com/snark-keys-ro.o1test.net}"
    curl "$base_s3_url/$tarname.tar.gz" | tar -xz -C "$tardir/web"

    for kind in {packaged,web,generated}; do
        $MINA_ROCKSDB_SCANNER --db-path "$tardir/$kind" --output-file "$WORKDIR/$kind.scan"
    done

    if ! ( cmp -s "$WORKDIR/generated.scan" "$WORKDIR/packaged.scan" && \
           cmp -s "$WORKDIR/packaged.scan" "$WORKDIR/web.scan" ); then
        echo "Error: kvdb contents mismatch for $tarname" >&2
        error=1
    fi
done

if [ $error -ne 0 ]; then
    echo "Error: failed validation" >&2
    exit 1
else
    echo "Validation successful" >&2
    exit 0
fi
